#   This program is free software; you can redistribute it and/or modify
#   it under the terms of the version 3 of the GNU Lesser General Public License
#   as published by the Free Software Foundation.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU Lesser General Public License
#   along with this program; if not, write to the Free Software
#   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
#
# Copyright (c) NEC Deutschland GmbH, NEC HPC Europe
#
# $Id$
import logging
from basic_types import *
from basic_types.factory import Factory
from basic_types.helpers import md5sum_file
from aggmon.aggmon_component import Component


__all__ = [ "Group", "Host", "Hierarchy", "HierarchyComponent" ]


class HNode(KeyValueObject):
    def __init__( self, name=None, parent_group=None ):
        self._name = name
        self._parent_group_name = parent_group

    def get_name( self ):
        return self._name
    
    def get_parent_group_name( self ):
        return self._parent_group_name


class Group(HNode):
    def __init__( self, name=None, parent_group=None ):
        HNode.__init__( self, name, parent_group )
        self._groups = Map()
        self._hosts = Map()
        self._master_host_names = Sequence()
        #print "Created GROUP", name, parent_group, self

    def add_group( self, group ):
        assert isinstance( group, Group ), "group is not a Group!"
        group_name = group.get_name()
        #print "Adding group", group_name, "to group", self.get_name()
        self._groups[group_name] = group

    def add_master_name( self, master_name ):
        self._master_host_names.append( master_name )

    def add_host( self, host ):
        assert isinstance( host, Host ), "host is not a Host!"
        host_name = host.get_name()
        #print "Adding host", host.get_name(), "to group", self.get_name()
        self._hosts[host_name] = host

    def dump( self ):
        logging.info( "Group name: %s" % self.get_name() )
        logging.info( "  parent_group name: %s" % self.get_parent_group_name() )
        logging.info( "  Hosts:\n" + 
                      "       " + "\n       ".join( host.get_name() for host in self.get_hosts()) )
        logging.info( "  Groups:\n" +
                      "       " + "\n       ".join( group for group in self.get_group_names()) )
        logging.info( "  Masters:\n" +
                      "       " + "\n       ".join( host for host in self.get_master_names()) )

    def get_group_names( self ):
        return self._groups.keys()

    def get_groups( self ):
        return self._groups.values()
    
    def get_host_names( self ):
        return self._hosts.keys()

    def get_hosts( self ):
        return self._hosts.values()
    
    def get_master_names( self ):
        return self._master_host_names

    def get_member_short_names( self ):
        members = [name.split("/")[-1] for name in self.get_host_names()]
        members.extend( [name.split("/")[-1] for name in self.get_group_names()] )
        return members


class Host(HNode):
    def __init__( self, name=None, parent_group=None ):
        HNode.__init__( self, name, parent_group )
        # child groups are the groups for which this host is master
        # A host can be master of his own group or of children of the own group
        # A host that is master of the own group can be also master of the parent of the own group
        self._child_group_names = Sequence()

    def add_child_group_name( self, group_name ):
        self._child_group_names.append( group_name )

    def get_child_group_names( self ):
        return self._child_group_names


Factory.set_factory( "Group", Group, __name__ )
Factory.set_factory( "Host", Host, __name__ )


class Hierarchy(object):
    """
    Class that presents the group/host hierarchy of cluster hosts.
    Hosts belong to groups, groups have one or more masters.
    The hierarchy is represented by a path to the host, for
    example: /europe/germany/kiel/rack2/node17
    Master hosts have tags like master:/europe/germany/kiel ,
    i.e. describing the hierarchy path of the group they are masters of.   
    """

    ROOT_NAME = "universe"
    ROOT_PATH = "/"

    def __init__( self, hostpath=None, config_file=None, config_string=None, md5sum=None, universe=None, verbose=False ):
        if universe is not None:
            self._clear_universe()
            self.__root = universe
        else:
            if config_file is None and config_string is not None:
                from cStringIO import StringIO
                config_file = StringIO( config_string )
            try:
                self.load_config_file( config_file, md5sum=md5sum, verbose=verbose )
            except:
                raise

        # validate hierarchy of masters
        self._is_valid( verbose=verbose )
        #self._is_cyclic( verbose=verbose )

        # locate own host in hierarchy, complain/warn if not found
        own_host = self.get_hosts( hostpath )
        if own_host is None:
            # if this is a hostname, check if there is exactly one host matching it
            match_num = 0
            match_path = ""
            for host in self.get_host_names():
                if host.endswith( "/" + hostpath ):
                    match_num = match_num + 1
                    match_path = host
            if match_num == 1:
                self._own_host = self.get_hosts( match_path )[0]
                if verbose:
                    logging.info( "Own host '%s' found in hierarchy at '%s'. Good." % (hostpath, self._own_host.get_name()) )
            else:
                print "WARNING: Own host '%s' not found in group hierarchy!" % hostpath
                print "         Are you sure you passed the full host hierarchy path?"
        else:
            if verbose:
                print "Own host '%s' found in hierarchy. Good." % hostpath
            self._own_host = own_host[0]

    def _add_group_master( self, host=None, attributes=[] ):
        for attribute in attributes:
            if not attribute.startswith( "m:/" ):
                continue
            group = self._add_groups_path( self.__root,
                                           attribute.lstrip( "m:/" ).split( "/" ) )
            host.add_child_group_name( group.get_name() )
            group.add_master_name( host.get_name() )
            #print "MASTER added", host.get_name(), "is master of group", group.get_name()

    def add_group( self, group=None ):
        assert group is not None, "none group added!?"
        group_name = group.get_name()
        self.__groups[group_name] = group

    def add_host( self, host=None ):
        assert host is not None, "none host added!?"
        assert isinstance( host, Host ), "host is not a Host!"
        host_name = host.get_name()
        self.__hosts[host_name] = host

    def _add_groups_path( self, parent_group=None, group_path_names=[] ):
        pgroup = parent_group
        group_name = ""
        for group_member_name in group_path_names:
            group = {}
            if len( group_member_name ) == 0:
                pgroup = self.__root
                continue
            group_name = group_name + "/" + group_member_name
            if group_name in self.__groups:
                pgroup = self.__groups[group_name]
                continue
            group = Group( group_name, pgroup.get_name() )
            self.add_group( group )
            pgroup.add_group( group )
            pgroup = group
        return pgroup
    
    def _clear_universe( self ):
        self.__hosts = {}
        self.__groups = {}
        self.__root = None
        self._own_host = None

    def get_groups( self, group_name=None ):
        """
        Return groups in hierarchy:
        Without argument: an array of all groups in the hierarchy.
        With group_name argument: a particular group.
        """
        if group_name is None:
            return self.__groups.values()
        elif group_name in self.__groups:
            return [self.__groups[group_name]]
        return None

    def get_group_masters( self, host ):
        """
        Get master(s) of the group the host belongs to. Info for the group is
        published there and consumers of group-specific events subscribe there.
        A host can be it's own master!
        """
        return [ self.get_hosts( host_name )[0] for host_name in
                    self.get_groups( host.get_parent_group_name() )[0].get_master_names() ]

    def get_group_names( self ):
        "Return all group names in the hierarchy."
        return self.__groups.keys()

    def get_own_host( self ):
        "Return own host object."
        return self._own_host

    def get_hosts( self, host_name=None ):
        """
        Return hosts in hierarchy:
        Without argument: an array of all hosts in the hierarchy.
        With host_name argument: a particular host.
        """
        if host_name is None:
            return self.__hosts.values()
        elif host_name in self.__hosts:
            return [self.__hosts[host_name]]
        return None

    def get_host_names( self ):
        "Return all host names in the hierarchy."
        return self.__hosts.keys()

    def get_superiors( self, host ):
        """
        Get "upstream" master(s). These are the hosts a group master would report to.
        """
        group = self.get_groups( host.get_parent_group_name() )
        if group is not self.__root:
            pgroup = self.get_groups( group.get_parent_group_name() )
            masters = [self.get_hosts( host_name ) for host_name in pgroup.get_master_names() ]  
            while masters.count( host ) > 0:
                masters = masters.remove( host )
                # may happen this was the only element, then the array is gone 
                if masters is None:
                    masters = []
                if pgroup.get_parent_group_name() is not None:
                    pgroup = self.get_groups( pgroup.get_parent_group_name() )
                    masters.extend( ( self.get_hosts( host_name ) 
                                      for host_name in pgroup.get_master_names() ) )

            return masters
        else:
            return None

    def _get_masters_or_superiors( self, host ):
        """
        Return a set of hosts to which this host will send events.
        This is either the group master, or (if the host is a group master)
        the superiors of the group.
        """
        masters = self.get_group_masters( host )
        #print "    masters:", [m.get_name() for m in masters]
        if masters.count( host ) > 0:
            masters = self.get_superiors( host )
        return masters

    def _get_depth( self, host, depth, max_depth, verbose=False ):
        """
        Function called recursively that determines the depth within the topology tree
        of a particular host.
          - host: the host for which we try to find the depth inside the hierarchy tree
          - depth: current depth, passed and increases during the recursion
          - max_depth: depth to stop at, if we reach this, we have a loop
        """
        if verbose:
            print " (%d) %s" % (depth, host.get_name())
        my_depth = depth + 1
        assert max_depth > my_depth, "Hierarchy loop detected at host %s" % host.get_name()
        masters = self._get_masters_or_superiors( host )
        if masters is not None:
            #print "masters:", masters
            old_depth = my_depth
            for master in masters:
                master_depth = self._get_depth( master, old_depth, max_depth, verbose=verbose )
                if master_depth > my_depth:
                    my_depth = master_depth
        return my_depth

    def in_own_groups( self, short_host_name ):
        groups = self.get_own_host().get_child_group_names()
        for group_name in groups:
            short_host_names = self.get_groups( group_name=group_name )[0].get_member_short_names()
            if short_host_name in short_host_names:
                return True
        return False

    def _is_cyclic( self, verbose=False ):
        for host_name, host in self.__hosts.items():
            if verbose:
                print "Checking host", host_name
            max_depth = len( host_name.split( "/" ) )
            depth = self._get_depth( host, 0, 2*max_depth+1, verbose=verbose )
            if verbose:
                print " %s depth=%d" % (host_name, depth)

    def _is_valid( self, verbose=False ):
        """
        Check validity of hierarchy.
        Currently limited to following conditions:
        - all groups exist
        - each group has at least one master (and the master is known)
        - group parents are correct
        - group universe has no parent group 
        """
        for host_name in self.__hosts.keys():
            group_names = host_name.split( "/" )[1:-1]
            gpath = ""
            for g in group_names:
                gpath = gpath + "/" + g
                assert gpath in self.__groups, \
                    "Hierarchy broken: missing group '%s' (from path of host '%s')" % (gpath, host_name)
            if len( gpath ) > 0:
                assert self.__hosts[host_name].get_parent_group_name() == gpath, \
                    "Hierarchy broken: host '%s' must have parent '%s'" % (host_name, gpath)
            else:
                # TODO: check if universe is the parent of this one
                pass

        # check that every defined group has masters
        for group in self.__groups.values():
            assert len( group.get_master_names() ) > 0, "Hierarchy broken: group '%s' has no masters!" % group.get_name()

        # each group has the parent it should have
        for group_name, group in self.__groups.items():
            if group_name == Hierarchy.ROOT_PATH:
                continue
            parent_group_path =  "/" + "/".join( group_name.split( "/" )[1:-1] )
            logging.debug( "group_name %s" % group_name )
            if parent_group_path != "/":
                assert group.get_parent_group_name() == parent_group_path, \
                    "Hierarchy broken: group '%s' has wrong parent: '%s' (should be '%s')" \
                    % (group_name, group.get_parent_group_name(), parent_group_path)
            else:
                assert group.get_parent_group_name() == Hierarchy.ROOT_PATH, \
                    "Hierarchy broken: group '%s' has wrong parent: '%s' (should be '%s')" \
                    % (group_name, group.get_parent_group_name(), Hierarchy.ROOT_PATH)

        # root group should have no parent
        assert self.__root.get_parent_group_name() is None, "Hierarchy broken: root group must not have a parent!"


    def load_config_file( self, config_file=None, md5sum=None, verbose=False ):
        self._clear_universe()
        # open, read line for line, interpret lines
        if config_file is None:
            return
        if isinstance( config_file, basestring ):
            try:
                f = open( config_file, "r" )
            except:
                logging.error( "Failed to open hierarchy file %s" % config_file )
                return
        else:
            # assume this is an object we can read from
            f = config_file
        try:
            if md5sum is not None:
                file_sum = md5sum_file( f )
                f.seek( 0 )
                if file_sum != md5sum:
                    logging.error( "md5sum of hierarchy file is wrong! Did the file change?" )
                    # TODO: What else do we do here? Fail? How?
                    # return
            lines = f.readlines()
        except:
            logging.error( "Could not read from (assumed) file-like object" )
            return
        try:
            f.close()
        except:
            pass
        # we usually call the root group "universe"
        self.__root = Group( Hierarchy.ROOT_PATH )
        self.add_group( self.__root )
        for line in lines:
            words = line.rstrip( "\n" ).split()
            if len( words ) < 1 or words[0].startswith( "#" ):
                continue
            host_path = words[0]
            assert host_path.startswith( "/" ), "Host path %s doesn't start with /" % host_path
            elements = host_path.split( "/" )[1:]
            parent_group = self._add_groups_path( self.__root, elements[0:-1] )
            host_name = host_path
            if host_name in self.__hosts:
                logging.warning( "duplicate host in hierarchy file! Line: %s" % line )
            else:
                host = Host( host_name, parent_group.get_name() )
                if len( words ) > 1:
                    self._add_group_master( host, words[1:] )
                self.add_host( host )
                parent_group.add_host( host )

    @staticmethod
    def path_to_topic( path ):
        "Convert a hierarchy path to a fragment suitable as topic addressing piece"
        return ".".join( path.split( "/" )[1:] )

    def print_info( self ):
        logging.info( "Hierarchy info:" )
        for group in self.__groups.values():
            group.dump()


class HierarchyComponent(Component):
    def __init__( self, *args, **kwds ):
        Component.__init__( self, *args, name="Hierarchy", **kwds )
        self.allowed_kwds = ["hostpath", "config_file", "config_string", "md5sum", "universe", "verbose"]

    def instantiate( self, *args, **kwds ):
        # lose the reference to any previous instance
        self.instance = None
        try:
            logging.debug( "instantiating: %s" % self.name )
            self.instance = Hierarchy( *args, **kwds )
        except Exception, e:
            raise e

    def start( self, **kwds ):
        "Equivalent to instantiating the hierarchy."
        k = {}
        for key, val in self.cmd["kwds"].items():
            k[key] = val 
        for key, val in kwds.items():
            if key in self.allowed_kwds:
                k[key] = val
        try:
            self.instantiate( *self.cmd["args"], **k )
        except Exception, e:
            self.state = Component.STOPPED
            logging.error( "Failed starting: %s, Exception: %s" % (self.name, str(e)) )
            raise e
        else:
            self.state = Component.RUNNING

    def stop( self ):
        # just forget the current instance?
        self.instance = None
        self.state = Component.STOPPED
